package lc2;
import java.util.*;
import org.junit.*;
public class Ex692 {
    @Test
    public void test() {
        Solution s = new Solution();
        String[] words = {"i", "love", "leetcode", "i", "love", "coding"};
        words = new String[]{"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"};
        // words = new String[]{"aaa","aa","a"};

        // System.out.println(s.topKFrequent(words, 1));
        System.out.println(s.topKFrequent(words, 4));
    }


    class Solution {
        public List<String> topKFrequent(String[] words, int k) {
            List<String> res = new ArrayList<>();
            Heap heap = new Heap(words.length);
            for (String s: words) {
                heap.insert(s);
            }
            for (int i = 0; i < k; i++) {
                res.add(heap.pop());
            }
            return res;
        }
    }

    class Heap {

        Node[] heap;
        int size;
        Map<String, Integer> idx;

        public Heap(int capacity) {
            heap = new Node[capacity];
            size = 0;
            idx = new HashMap<>();
        }

        public void insert(String s) {
            if (!idx.containsKey(s)) {
                Node node = new Node(s);
                idx.put(s, size);
                heap[size] = node;
                adjust(node, size++);
            } else {
                int i = idx.get(s);
                Node node = heap[i];
                node.increase();
                adjust(node, i);
            }
        }

        private void adjust(Node node, int i) { 
            while (i > 0) {
                int p = (i - 1) / 2;
                if (node.compareTo(heap[p]) > 0) {
                    swap(i, p);
                    i = p;
                } else {
                    break;
                }
            }
        }

        private void swap(int i, int j) {
            String si = heap[i].word, sj = heap[j].word;
            int v = idx.get(si);
            idx.put(si, idx.get(sj));
            idx.put(sj, v);
            Node t = heap[i];
            heap[i] = heap[j];
            heap[j] = t;
        }

        public String pop() {
            String s = heap[0].word;
            System.out.println(s + ": " + heap[0].times);
            swap(0, --size);
            adjustDown();
            return s;
        }

        private void adjustDown() {
            int i = 0;
            while (i < size) {
                int l = i * 2 + 1, r = i * 2 + 2;
                int max = i;
                if (l < size) {
                    if (heap[max].compareTo(heap[l]) < 0) {
                        max = l;
                    }
                }
                if (r < size) {
                    if (heap[max].compareTo(heap[r]) < 0) {
                        max = r;
                    }
                }
                if (max == i) {
                    break;
                } else {
                    swap(i, max);
                    i = max;
                }
            }
        }
    }

    class Node implements Comparable<Node> {
        String word;
        int times;

        public Node(String s) {
            this.word = s;
            this.times = 1;
        }

        public void increase() {
            this.times++;
        }

        @Override
        public int compareTo(Node o) {
            if (this.times != o.times) {
                return this.times - o.times;
            } else {
                return -this.word.compareTo(o.word);
            }
        }
    }
}
